import os
import cv2
import sys
import torch
import numpy as np

from ultralytics import YOLO
from PyQt6 import QtCore, QtGui, QtWidgets
from PyQt6.QtGui import QImage, QPixmap
from PyQt6.QtCore import Qt

from interface.MainWindow import Ui_MainWindow


def main():
    """程序入口点"""
    Application()


class Application:
    def __init__(self):
        app = QtWidgets.QApplication(sys.argv)

        # 配置全集变量
        self.cap = None  # cv2 的摄像头对象
        self.video = None  # cv2 视频对象
        self.model = None  # 加载的模型
        self.use_cuda = True  # 是否使用 cuda 计算

        # 主界面
        self.MainWindow = QtWidgets.QMainWindow()
        self.ui = Ui_MainWindow()

        # 组件配置
        self.timer_video = QtCore.QTimer()  # OpenCV 视频画面截取
        self.timer_cap = QtCore.QTimer()  # OpenCV 摄像头画面截取
        self.ui.setupUi(self.MainWindow)

        # 调用初始化函数
        self.init_MainWindow()

        # 显示
        self.MainWindow.show()
        sys.exit(app.exec())

    def init_MainWindow(self):
        """初始化主界面"""
        # 绑定事件
        self.MainWindow.closeEvent = self.exit_application  # 右上角的 ×
        self.ui.cuda_checkBox.clicked.connect(lambda: self.set_option("cuda"))  # 是否使用 cuda
        self.ui.exitBtn.clicked.connect(self.exit_application)  # 关闭按钮
        self.ui.selectModel.clicked.connect(self.load_model)  # 选择模型
        self.ui.openVideoBtn.clicked.connect(self.start_video)  # 选择视频
        self.ui.stopDetectBtn.clicked.connect(self.stop_detect)  # 停止视频
        self.ui.openCamBtn.clicked.connect(self.start_camera)  # 开始摄像头
        self.timer_video.timeout.connect(self.video_show)
        self.timer_cap.timeout.connect(self.camera_show)

    def load_model(self):
        fileName, fileType = QtWidgets.QFileDialog.getOpenFileName(self.MainWindow, "选取模型权重", filter='*.pt')

        if fileName.endswith('.pt'):
            self.release_model()  # 已存在则先释放
            if self.use_cuda:
                self.model = YOLO(fileName).cuda()
            else:
                self.model = YOLO(fileName)

            self.ui.openVideoBtn.setEnabled(True)
            self.ui.stopDetectBtn.setEnabled(True)
            self.ui.openCamBtn.setEnabled(True)

            print(f"模型 {fileName} 已经成功加载。")
        else:
            QtWidgets.QMessageBox.information(self.MainWindow, "选择模型不正确", "请选择以正确的模型。")

    def detect_result(self, frame, conf, label, position):
        """检测结果回调"""
        # 例如：
        # frame: CV frame
        # conf: 0 ~ 1
        # label: person
        # position: [355.4978332519531, 33.8280143737793, 376.4524230957031, 98.86994934082031]
        pass

    def start_camera(self):
        """启动摄像头"""
        self.timer_cap.stop()
        self.timer_video.stop()

        if self.cap is None:
            self.cap = cv2.VideoCapture(int(self.ui.camera_id_spinBox.text()), cv2.CAP_DSHOW)
        if self.cap.isOpened():
            self.timer_cap.start(50)
        self.ui.stopDetectBtn.setEnabled(True)

    def start_video(self):
        """启动视频"""
        self.timer_cap.stop()
        self.timer_video.stop()

        fileName, fileType = QtWidgets.QFileDialog.getOpenFileName(self.MainWindow, "选取视频文件", filter='*.mp4')
        if os.path.isfile(fileName):
            self.video = cv2.VideoCapture(fileName)
            fps = self.video.get(cv2.CAP_PROP_FPS)
            self.timer_video.start(int(1000 / fps))
        else:
            QtWidgets.QMessageBox.information(self.MainWindow, "选择视频文件不正确", "请选择正确的文件。")

    def video_show(self):
        """展示一帧"""
        ret, frame = self.video.read()
        if ret:
            self.process_frame(frame)
        else:
            self.timer_video.stop()
            img = cv2.cvtColor(np.zeros((500, 500), np.uint8), cv2.COLOR_BGR2RGB)
            img = QtGui.QImage(img.data, img.shape[1], img.shape[0], QtGui.QImage.Format.Format_RGB888)
            self.ui.oriVideoLabel.setPixmap(QtGui.QPixmap.fromImage(img))
            self.ui.detectlabel.setPixmap(QtGui.QPixmap.fromImage(img))
            self.video.release()
            self.video = None

    def camera_show(self):
        ret, frame = self.cap.read()
        if ret:
            self.process_frame(frame)

    def process_frame(self, frame):
        if self.model is not None:
            # 保留原始尺寸的图像
            original_frame = frame.copy()

            # 进行目标检测（不缩放）
            device = 'cuda' if torch.cuda.is_available() and self.use_cuda else 'cpu'
            results = self.model(frame, device=device, verbose=False)  # 使用原始大小的图像进行检测

            # 提取检测结果
            for result in results:
                boxes = result.boxes  # 获取所有检测框
                for box in boxes:
                    # 获取框的位置
                    x1, y1, x2, y2 = box.xyxy[0].tolist()  # 左上角(x1, y1)，右下角(x2, y2)

                    # 获取框的类别和置信度
                    conf = box.conf.item()  # 置信度
                    cls = int(box.cls.item())  # 类别索引
                    label = self.model.names[cls]  # 类别名称

                    self.detect_result(frame, conf, label, (x1, y1, x2, y2))

        detected_frame = results[0].plot()

        # 缩放检测后的图像
        resized_detected_frame = cv2.resize(detected_frame, (448, 352))
        resized_detected_frame = cv2.cvtColor(resized_detected_frame, cv2.COLOR_RGB2BGR)
        qimage_detected = QImage(resized_detected_frame.data, resized_detected_frame.shape[1],
                                 resized_detected_frame.shape[0], resized_detected_frame.strides[0],
                                 QImage.Format.Format_RGB888)
        pixmap_detected = QPixmap.fromImage(qimage_detected)
        scaled_pixmap_detected = pixmap_detected.scaled(self.ui.detectlabel.size(), Qt.AspectRatioMode.KeepAspectRatio)
        self.ui.detectlabel.setAlignment(Qt.AlignmentFlag.AlignCenter)  # 设置 QLabel 的对齐方式为居中
        self.ui.detectlabel.setPixmap(scaled_pixmap_detected)

        # 缩放并显示原始图像
        resized_original_frame = cv2.resize(original_frame, (448, 352))
        resized_original_frame = cv2.cvtColor(resized_original_frame, cv2.COLOR_BGR2RGB)
        qimage_original = QImage(resized_original_frame.data, resized_original_frame.shape[1],
                                 resized_original_frame.shape[0], resized_original_frame.strides[0],
                                 QImage.Format.Format_RGB888)
        pixmap_original = QPixmap.fromImage(qimage_original)
        scaled_pixmap_original = pixmap_original.scaled(self.ui.oriVideoLabel.size(),
                                                        Qt.AspectRatioMode.KeepAspectRatio)
        self.ui.oriVideoLabel.setAlignment(Qt.AlignmentFlag.AlignCenter)  # 设置 QLabel 的对齐方式为居中
        self.ui.oriVideoLabel.setPixmap(scaled_pixmap_original)

    def set_option(self, signal):
        """设置"""
        if signal == "cuda":
            self.use_cuda = self.ui.cuda_checkBox.checkState()

    def stop_detect(self):
        if self.timer_video.isActive():
            self.timer_video.stop()
        if self.timer_cap.isActive():
            self.timer_cap.stop()
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        self.video = None
        img = cv2.cvtColor(np.zeros((500, 500), np.uint8), cv2.COLOR_BGR2RGB)
        img = QtGui.QImage(img.data, img.shape[1], img.shape[0], QtGui.QImage.Format.Format_RGB888)
        self.ui.oriVideoLabel.setPixmap(QtGui.QPixmap.fromImage(img))
        self.ui.detectlabel.setPixmap(QtGui.QPixmap.fromImage(img))

    def release_model(self):
        """释放当前加载的模型"""
        if self.model is not None:
            del self.model  # 删除模型对象
            self.model = None
            if self.use_cuda:
                torch.cuda.empty_cache()  # 清空未使用的显存

    def exit_application(self, event=0):
        """程序退出时"""
        if self.cap is not None:
            self.cap.release()
            self.cap = None
        if self.cap is not None:
            self.video.release()
            self.video = None
        if self.timer_cap.isActive():
            self.timer_cap.stop()
        if self.timer_video.isActive():
            self.timer_video.stop()
        exit()


if __name__ == "__main__":
    main()
